# utils/helpers.py
import os
import torch
import logging


def save_model(model, model_dir, epoch, best_epoch=None):
    """
    保存模型检查点，并可选择性地删除旧的最佳模型。
    :param model: 要保存的模型。
    :param model_dir: 模型保存的目录。
    :param epoch: 当前训练轮次。
    :param best_epoch: (可选) 上一个最佳轮次，用于清理旧模型。
    """
    if not os.path.exists(model_dir):
        os.makedirs(model_dir)

    # 定义模型文件名
    model_state_file = os.path.join(model_dir, f'model_epoch_{epoch}.pth')

    # 为了兼容性和可移植性，建议保存到CPU上
    state = {
        'state_dict': model.state_dict(),
        'epoch': epoch,
    }
    torch.save(state, model_state_file)
    logging.info(f"模型已保存到: {model_state_file}")

    # 清理旧的最佳模型文件
    if best_epoch is not None and epoch != best_epoch:
        old_model_state_file = os.path.join(model_dir, f'model_epoch_{best_epoch}.pth')
        if os.path.exists(old_model_state_file):
            try:
                os.remove(old_model_state_file)
                logging.info(f"已删除旧的最佳模型: {old_model_state_file}")
            except OSError as e:
                logging.error(f"删除旧模型失败: {e}")


def load_model(model, model_path, device):
    """
    加载模型状态。
    :param model: 要加载状态的目标模型。
    :param model_path: 模型文件的路径。
    :param device: 要将模型加载到的设备。
    :return: 加载状态后的模型。
    """
    logger = logging.getLogger()
    if model_path and os.path.exists(model_path):
        try:
            checkpoint = torch.load(model_path, map_location=device)
            model.load_state_dict(checkpoint['state_dict'])
            model.to(device)
            logger.info(f"模型从 {model_path} 加载成功，并已移动到 {device}。")
            model.eval()  # 加载后默认为评估模式
        except Exception as e:
            logger.error(f"加载模型失败: {e}")
    else:
        logger.warning(f"预训练模型文件未找到: {model_path}。将从头开始训练。")
    return model
